#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Panel图像拼接脚本
使用SAM模型提取对象和背景，进行图像拼接并更新forgery.json
"""

import os
import sys
import cv2
import numpy as np
import torch
import json
import argparse
import glob
import shutil
from pathlib import Path
from typing import Tuple, List, Dict, Any
import gc
from itertools import combinations
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction import image as skimage
import random

# 添加SAM模型路径
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'base_scripts', 'models', 'segment-anything-main'))
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry


class PanelSplicingProcessor:
    """Panel图像拼接处理器"""
    
    def __init__(self, device: str = "cuda", model_path: str = None):
        """
        初始化处理器
        
        Args:
            device: 计算设备
            model_path: SAM模型路径
        """
        self.device = device
        self.model_path = model_path or "/home/zhangbo/workspace/aigc/SAM/vit_models/sam_vit_h_4b8939.pth"
        self.sam_model = None
        self.mask_generator = None
        
    def load_sam_model(self):
        """加载SAM模型"""
        if self.sam_model is None:
            print("=" * 60)
            print("🔄 阶段1: 正在加载SAM模型...")
            print(f"📁 模型路径: {self.model_path}")
            print(f"🖥️  计算设备: {self.device}")
            print("=" * 60)
            
            model_type = "vit_h"
            print(f"📦 模型类型: {model_type}")
            print("⏳ 正在从检查点加载模型...")
            
            self.sam_model = sam_model_registry[model_type](checkpoint=self.model_path)
            self.sam_model.to(device=self.device)
            
            print("⚙️  正在配置掩码生成器...")
            # 配置掩码生成器
            self.mask_generator = SamAutomaticMaskGenerator(
                self.sam_model,
                output_mode="binary_mask",
                points_per_side=16,
                pred_iou_thresh=0.86,
                stability_score_thresh=0.92,
                min_mask_region_area=100
            )
            print("✅ SAM模型加载完成!")
            print("=" * 60)
    
    def calculate_panel_similarity(self, img1: np.ndarray, img2: np.ndarray) -> float:
        """
        计算两个panel图像的相似度
        
        Args:
            img1: 第一个图像
            img2: 第二个图像
            
        Returns:
            相似度分数 (0-1)
        """
        print(f"    🔍 正在计算图像相似度...")
        print(f"    📏 图像1尺寸: {img1.shape[1]}x{img1.shape[0]}")
        print(f"    📏 图像2尺寸: {img2.shape[1]}x{img2.shape[0]}")
        
        # 调整图像大小到相同尺寸
        target_size = (224, 224)
        print(f"    🔄 调整图像尺寸到: {target_size}")
        img1_resized = cv2.resize(img1, target_size)
        img2_resized = cv2.resize(img2, target_size)
        
        # 转换为灰度图
        print(f"    🎨 转换为灰度图...")
        gray1 = cv2.cvtColor(img1_resized, cv2.COLOR_BGR2GRAY)
        gray2 = cv2.cvtColor(img2_resized, cv2.COLOR_BGR2GRAY)
        
        # 计算直方图
        print(f"    📊 计算直方图...")
        hist1 = cv2.calcHist([gray1], [0], None, [256], [0, 256])
        hist2 = cv2.calcHist([gray2], [0], None, [256], [0, 256])
        
        # 计算直方图相关性
        correlation = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
        print(f"    📈 直方图相关性: {correlation:.3f}")
        
        # 计算结构相似性
        ssim_score = self.calculate_ssim(gray1, gray2)
        print(f"    🏗️  结构相似性(SSIM): {ssim_score:.3f}")
        
        # 计算尺寸相似性
        size_similarity = 1.0 - abs(img1.shape[0] - img2.shape[0]) / max(img1.shape[0], img2.shape[0])
        size_similarity *= 1.0 - abs(img1.shape[1] - img2.shape[1]) / max(img1.shape[1], img2.shape[1])
        print(f"    📐 尺寸相似性: {size_similarity:.3f}")
        
        # 综合相似度
        overall_similarity = (correlation * 0.4 + ssim_score * 0.4 + size_similarity * 0.2)
        print(f"    🎯 综合相似度: {overall_similarity:.3f}")
        
        return max(0, min(1, overall_similarity))
    
    def calculate_ssim(self, img1: np.ndarray, img2: np.ndarray) -> float:
        """计算结构相似性指数"""
        # 简化的SSIM计算
        mu1 = np.mean(img1)
        mu2 = np.mean(img2)
        sigma1 = np.var(img1)
        sigma2 = np.var(img2)
        sigma12 = np.mean((img1 - mu1) * (img2 - mu2))
        
        c1 = 0.01 ** 2
        c2 = 0.03 ** 2
        
        ssim = ((2 * mu1 * mu2 + c1) * (2 * sigma12 + c2)) / \
               ((mu1 ** 2 + mu2 ** 2 + c1) * (sigma1 + sigma2 + c2))
        
        return ssim
    
    def select_similar_panels(self, panel_dir: str, method: str = "similarity") -> Tuple[str, str]:
        """
        选择两个相似的panel图像
        
        Args:
            panel_dir: panel图像目录
            method: 选择方法 ("similarity" 或 "random")
            
        Returns:
            选中的两个panel文件路径
        """
        print("=" * 60)
        print("🔄 阶段2: 选择相似的Panel图像")
        print(f"📁 搜索目录: {panel_dir}")
        print(f"🎯 选择方法: {method}")
        print("=" * 60)
        
        panel_files = glob.glob(os.path.join(panel_dir, "*.png"))
        print(f"📊 找到 {len(panel_files)} 个PNG文件")
        
        if len(panel_files) < 2:
            raise ValueError(f"❌ panel目录中图像数量不足，需要至少2个，实际有{len(panel_files)}个")
        
        if method == "random":
            # 随机选择两个panel
            print("🎲 使用随机选择方法...")
            selected = random.sample(panel_files, 2)
            print(f"✅ 随机选择了panel: {os.path.basename(selected[0])} 和 {os.path.basename(selected[1])}")
            return selected[0], selected[1]
        
        elif method == "similarity":
            # 基于相似度选择
            print("🔍 使用相似度选择方法...")
            best_similarity = 0
            best_pair = None
            total_combinations = len(list(combinations(panel_files, 2)))
            
            print(f"📈 需要比较 {total_combinations} 对图像")
            print("-" * 40)
            
            for i, (img1_path, img2_path) in enumerate(combinations(panel_files, 2), 1):
                print(f"🔄 比较 {i}/{total_combinations}: {os.path.basename(img1_path)} vs {os.path.basename(img2_path)}")
                
                img1 = cv2.imread(img1_path)
                img2 = cv2.imread(img2_path)
                
                if img1 is None or img2 is None:
                    print(f"    ⚠️  跳过无效图像")
                    continue
                
                similarity = self.calculate_panel_similarity(img1, img2)
                print(f"    🎯 相似度: {similarity:.3f}")
                
                if similarity > best_similarity:
                    best_similarity = similarity
                    best_pair = (img1_path, img2_path)
                    print(f"    🏆 新的最佳匹配!")
                
                print("-" * 40)
            
            if best_pair is None:
                raise ValueError("❌ 无法找到有效的panel对")
            
            print("=" * 60)
            print(f"✅ 选择相似度最高的panel对:")
            print(f"   📄 Panel A: {os.path.basename(best_pair[0])}")
            print(f"   📄 Panel B: {os.path.basename(best_pair[1])}")
            print(f"   🎯 相似度: {best_similarity:.3f}")
            print("=" * 60)
            return best_pair[0], best_pair[1]
        
        else:
            raise ValueError(f"❌ 不支持的选择方法: {method}")
    
    def extract_object_and_background(self, image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        使用SAM模型提取对象和背景
        
        Args:
            image: 输入图像
            
        Returns:
            (object_mask, background_mask)
        """
        print("    🔄 正在提取对象和背景...")
        print(f"    📏 图像尺寸: {image.shape[1]}x{image.shape[0]}")
        
        if self.mask_generator is None:
            print("    ⚙️  首次使用，正在加载SAM模型...")
            self.load_sam_model()
        
        # 清理CUDA缓存
        print("    🧹 清理CUDA缓存...")
        torch.cuda.empty_cache()
        
        try:
            print("    🤖 正在使用SAM模型生成掩码...")
            # 生成masks
            masks = self.mask_generator.generate(image)
            
            if not masks:
                print("    ⚠️  警告: 未检测到任何对象，使用全图作为对象")
                object_mask = np.ones(image.shape[:2], dtype=np.uint8) * 255
                background_mask = np.zeros(image.shape[:2], dtype=np.uint8)
                return object_mask, background_mask
            
            print(f"    📊 检测到 {len(masks)} 个对象区域")
            
            # 按面积排序masks
            masks.sort(key=lambda x: x["area"], reverse=True)
            
            # 选择面积最大的mask作为对象
            largest_area = masks[0]["area"]
            print(f"    🎯 选择最大区域作为对象，面积: {largest_area} 像素")
            
            object_mask = masks[0]["segmentation"].astype(np.uint8) * 255
            
            # 背景是对象mask的补集
            background_mask = (1 - masks[0]["segmentation"]).astype(np.uint8) * 255
            
            print("    ✅ 对象和背景提取完成")
            return object_mask, background_mask
            
        except Exception as e:
            print(f"    ❌ 提取对象和背景时出错: {e}")
            print("    🔄 使用默认掩码...")
            # 返回默认mask
            object_mask = np.ones(image.shape[:2], dtype=np.uint8) * 255
            background_mask = np.zeros(image.shape[:2], dtype=np.uint8)
            return object_mask, background_mask
        finally:
            print("    🧹 清理内存...")
            torch.cuda.empty_cache()
            gc.collect()
    
    def resize_mask_to_image(self, mask: np.ndarray, target_image: np.ndarray) -> np.ndarray:
        """
        将mask调整到目标图像尺寸
        
        Args:
            mask: 原始mask
            target_image: 目标图像
            
        Returns:
            调整后的mask
        """
        target_height, target_width = target_image.shape[:2]
        resized_mask = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
        return resized_mask
    
    def create_synthetic_image(self, object_img: np.ndarray, object_mask: np.ndarray, 
                             background_img: np.ndarray, background_mask: np.ndarray) -> np.ndarray:
        """
        创建合成图像
        
        Args:
            object_img: 对象图像
            object_mask: 对象mask
            background_img: 背景图像
            background_mask: 背景mask
            
        Returns:
            合成后的图像
        """
        # 确保所有图像和mask具有相同的尺寸
        target_height = max(object_img.shape[0], background_img.shape[0])
        target_width = max(object_img.shape[1], background_img.shape[1])
        
        # 调整图像尺寸
        if object_img.shape[:2] != (target_height, target_width):
            object_img = cv2.resize(object_img, (target_width, target_height))
        if background_img.shape[:2] != (target_height, target_width):
            background_img = cv2.resize(background_img, (target_width, target_height))
        
        # 调整mask尺寸
        if object_mask.shape[:2] != (target_height, target_width):
            object_mask = cv2.resize(object_mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
        if background_mask.shape[:2] != (target_height, target_width):
            background_mask = cv2.resize(background_mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
        
        # 确保mask是3通道
        if len(object_mask.shape) == 2:
            object_mask = np.stack([object_mask] * 3, axis=-1)
        if len(background_mask.shape) == 2:
            background_mask = np.stack([background_mask] * 3, axis=-1)
        
        # 归一化mask到0-1范围
        object_mask_norm = object_mask.astype(np.float32) / 255.0
        background_mask_norm = background_mask.astype(np.float32) / 255.0
        
        # 合成图像
        synthetic_img = (object_img * object_mask_norm + 
                        background_img * background_mask_norm).astype(np.uint8)
        
        return synthetic_img
    
    def process_panel_splicing(self, panel_dir: str, output_dir: str, 
                             mask_output_dir: str, method: str = "similarity") -> Dict[str, Any]:
        """
        处理panel拼接任务
        
        Args:
            panel_dir: panel图像目录
            output_dir: 输出目录
            mask_output_dir: mask输出目录
            method: 选择方法
            
        Returns:
            处理结果信息
        """
        print("=" * 60)
        print("🚀 开始Panel拼接处理流程")
        print("=" * 60)
        
        # 创建输出目录
        print("📁 创建输出目录...")
        os.makedirs(output_dir, exist_ok=True)
        os.makedirs(mask_output_dir, exist_ok=True)
        print(f"   ✅ 合成图像目录: {output_dir}")
        print(f"   ✅ Mask图像目录: {mask_output_dir}")
        
        # 选择两个相似的panel
        panel_a_path, panel_b_path = self.select_similar_panels(panel_dir, method)
        
        # 读取图像
        print("=" * 60)
        print("🔄 阶段3: 读取和处理图像")
        print("=" * 60)
        
        print("📖 正在读取panel图像...")
        panel_a = cv2.imread(panel_a_path)
        panel_b = cv2.imread(panel_b_path)
        
        if panel_a is None or panel_b is None:
            raise ValueError("❌ 无法读取panel图像")
        
        print(f"✅ Panel A: {os.path.basename(panel_a_path)} - 尺寸: {panel_a.shape[1]}x{panel_a.shape[0]}")
        print(f"✅ Panel B: {os.path.basename(panel_b_path)} - 尺寸: {panel_b.shape[1]}x{panel_b.shape[0]}")
        
        # 提取panel A的对象和背景
        print("=" * 60)
        print("🔄 阶段4: 提取Panel A的对象和背景")
        print("=" * 60)
        obj_a_mask, bg_a_mask = self.extract_object_and_background(panel_a)
        
        # 提取panel B的对象和背景
        print("=" * 60)
        print("🔄 阶段5: 提取Panel B的对象和背景")
        print("=" * 60)
        obj_b_mask, bg_b_mask = self.extract_object_and_background(panel_b)
        
        # 调整mask尺寸以匹配目标图像
        print("=" * 60)
        print("🔄 阶段6: 调整Mask尺寸")
        print("=" * 60)
        print("📏 正在调整mask尺寸以匹配目标图像...")
        
        obj_a_mask_resized = self.resize_mask_to_image(obj_a_mask, panel_b)
        bg_a_mask_resized = self.resize_mask_to_image(bg_a_mask, panel_b)
        obj_b_mask_resized = self.resize_mask_to_image(obj_b_mask, panel_a)
        bg_b_mask_resized = self.resize_mask_to_image(bg_b_mask, panel_a)
        
        print(f"   ✅ A对象mask调整到B尺寸: {obj_a_mask_resized.shape[1]}x{obj_a_mask_resized.shape[0]}")
        print(f"   ✅ A背景mask调整到B尺寸: {bg_a_mask_resized.shape[1]}x{bg_a_mask_resized.shape[0]}")
        print(f"   ✅ B对象mask调整到A尺寸: {obj_b_mask_resized.shape[1]}x{obj_b_mask_resized.shape[0]}")
        print(f"   ✅ B背景mask调整到A尺寸: {bg_b_mask_resized.shape[1]}x{bg_b_mask_resized.shape[0]}")
        
        # 创建合成图像
        print("=" * 60)
        print("🔄 阶段7: 创建合成图像")
        print("=" * 60)
        print("🎨 正在创建合成图像...")
        
        # 合成图像C: A的对象 + B的背景
        print("   🔄 合成图像C: A的对象 + B的背景")
        synthetic_c = self.create_synthetic_image(panel_a, obj_a_mask_resized, panel_b, bg_b_mask_resized)
        
        # 合成图像D: B的对象 + A的背景
        print("   🔄 合成图像D: B的对象 + A的背景")
        synthetic_d = self.create_synthetic_image(panel_b, obj_b_mask_resized, panel_a, bg_a_mask_resized)
        
        print("   ✅ 合成图像创建完成")
        
        # 保存合成图像
        print("=" * 60)
        print("🔄 阶段8: 保存结果文件")
        print("=" * 60)
        
        panel_a_name = os.path.splitext(os.path.basename(panel_a_path))[0]
        panel_b_name = os.path.splitext(os.path.basename(panel_b_path))[0]
        
        synthetic_c_path = os.path.join(output_dir, f"{panel_a_name}_spliced.png")
        synthetic_d_path = os.path.join(output_dir, f"{panel_b_name}_spliced.png")
        
        print("💾 正在保存合成图像...")
        cv2.imwrite(synthetic_c_path, synthetic_c)
        cv2.imwrite(synthetic_d_path, synthetic_d)
        print(f"   ✅ 合成图像C: {synthetic_c_path}")
        print(f"   ✅ 合成图像D: {synthetic_d_path}")
        
        # 保存mask图像
        mask_a_path = os.path.join(mask_output_dir, f"{panel_a_name}_mask.png")
        mask_b_path = os.path.join(mask_output_dir, f"{panel_b_name}_mask.png")
        
        print("💾 正在保存mask图像...")
        cv2.imwrite(mask_a_path, obj_a_mask_resized)
        cv2.imwrite(mask_b_path, obj_b_mask_resized)
        print(f"   ✅ Mask A: {mask_a_path}")
        print(f"   ✅ Mask B: {mask_b_path}")
        
        # 返回处理结果信息
        result = {
            "original_panels": {
                "panel_a": panel_a_path,
                "panel_b": panel_b_path
            },
            "synthetic_panels": {
                "synthetic_c": synthetic_c_path,
                "synthetic_d": synthetic_d_path
            },
            "masks": {
                "mask_a": mask_a_path,
                "mask_b": mask_b_path
            },
            "panel_info": {
                "panel_a_name": panel_a_name,
                "panel_b_name": panel_b_name
            }
        }
        
        print("=" * 60)
        print("🎉 Panel拼接处理完成!")
        print("=" * 60)
        
        return result
    
    def update_forgery_json(self, forgery_json_path: str, result: Dict[str, Any], 
                           output_base_path: str = "dataset/forgery/forgery_pdf/2010072670_FG_001") -> None:
        """
        更新forgery.json文件
        
        Args:
            forgery_json_path: forgery.json文件路径
            result: 处理结果
            output_base_path: 输出基础路径
        """
        print("=" * 60)
        print("🔄 阶段9: 更新forgery.json文件")
        print("=" * 60)
        
        # 读取现有的forgery.json
        print(f"📖 正在读取forgery.json: {forgery_json_path}")
        if os.path.exists(forgery_json_path):
            with open(forgery_json_path, 'r', encoding='utf-8') as f:
                forgery_data = json.load(f)
            print(f"   ✅ 成功读取现有文件，包含 {len(forgery_data.get('panels', []))} 个panel条目")
        else:
            forgery_data = {"panels": []}
            print("   📝 文件不存在，创建新的数据结构")
        
        # 创建新的panel条目
        panel_a_name = result["panel_info"]["panel_a_name"]
        panel_b_name = result["panel_info"]["panel_b_name"]
        
        print(f"📝 正在创建panel条目...")
        print(f"   Panel A: {panel_a_name}")
        print(f"   Panel B: {panel_b_name}")
        
        # 为panel A创建条目
        panel_a_entry = {
            "panel_id": f"{panel_a_name}_FG_001",
            "bbox_page_pixel": [0, 0, 0, 0],  # 需要根据实际情况调整
            "path": f"{output_base_path}/panel/{panel_a_name}/{panel_a_name}_spliced.png",
            "panel_level_issues": {
                "has_issue": True,
                "issues": [
                    {
                        "issue_id": "FG_001",
                        "scope": "within",
                        "issue_type": "forgery",
                        "issue_subtype": "splicing",
                        "evidence": {
                            "mask_path": f"{output_base_path}/mask/{panel_a_name}/{panel_a_name}_mask.png"
                        }
                    }
                ]
            }
        }
        
        # 为panel B创建条目
        panel_b_entry = {
            "panel_id": f"{panel_b_name}_FG_001",
            "bbox_page_pixel": [0, 0, 0, 0],  # 需要根据实际情况调整
            "path": f"{output_base_path}/panel/{panel_b_name}/{panel_b_name}_spliced.png",
            "panel_level_issues": {
                "has_issue": True,
                "issues": [
                    {
                        "issue_id": "FG_001",
                        "scope": "within",
                        "issue_type": "forgery",
                        "issue_subtype": "splicing",
                        "evidence": {
                            "mask_path": f"{output_base_path}/mask/{panel_b_name}/{panel_b_name}_mask.png"
                        }
                    }
                ]
            }
        }
        
        print("   ✅ Panel条目创建完成")
        
        # 添加新条目到forgery数据
        print("📋 正在添加新条目到forgery数据...")
        forgery_data["panels"].extend([panel_a_entry, panel_b_entry])
        print(f"   ✅ 添加了2个新panel条目，总计: {len(forgery_data['panels'])} 个条目")
        
        # 保存更新后的forgery.json
        print("💾 正在保存更新后的forgery.json...")
        with open(forgery_json_path, 'w', encoding='utf-8') as f:
            json.dump(forgery_data, f, indent=2, ensure_ascii=False)
        
        print(f"✅ 已更新forgery.json: {forgery_json_path}")
        print("=" * 60)


def main():
    """主函数"""
    print("🎬 启动Panel图像拼接处理程序")
    print("=" * 80)
    parser = argparse.ArgumentParser(description='Panel图像拼接处理')
    parser.add_argument('--panel_dir', type=str,
                       default='/home/zhangbo/workspace/aigc/psy/splicing/2010072670_FG_004/panel/2010072670_FG_004_5_1',
                       help='panel图像目录')
    parser.add_argument('--output_dir', type=str,
                       default='/home/zhangbo/workspace/aigc/psy/splicing/2010072670_FG_004/synthetic_panels',
                       help='合成图像输出目录')
    parser.add_argument('--mask_output_dir', type=str,
                       default='/home/zhangbo/workspace/aigc/psy/splicing/2010072670_FG_004/masks',
                       help='mask图像输出目录')
    parser.add_argument('--forgery_json', type=str,
                       default='/home/zhangbo/workspace/aigc/psy/splicing/forgery.json',
                       help='forgery.json文件路径')
    parser.add_argument('--method', type=str, choices=['similarity', 'random'],
                       default='similarity', help='panel选择方法')
    parser.add_argument('--device', type=str, default='cuda', help='计算设备')
    parser.add_argument('--model_path', type=str,
                       default='/home/zhangbo/workspace/aigc/SAM/vit_models/sam_vit_h_4b8939.pth',
                       help='SAM模型路径')
    # parser = argparse.ArgumentParser(description='Panel图像拼接处理')
    # parser.add_argument('--panel_dir', type=str,
    #                    default='/Users/pppsy/code/python/ForensicsGenome/code/splicing/2010072670/panel/2010072670_5_1',
    #                    help='panel图像目录')
    # parser.add_argument('--output_dir', type=str,
    #                    default='/Users/pppsy/code/python/ForensicsGenome/code/splicing/2010072670/synthetic_panels',
    #                    help='合成图像输出目录')
    # parser.add_argument('--mask_output_dir', type=str,
    #                    default='/Users/pppsy/code/python/ForensicsGenome/code/splicing/2010072670/masks',
    #                    help='mask图像输出目录')
    # parser.add_argument('--forgery_json', type=str,
    #                    default='/Users/pppsy/code/python/ForensicsGenome/dataset/forgery/forgery.json',
    #                    help='forgery.json文件路径')
    # parser.add_argument('--method', type=str, choices=['similarity', 'random'],
    #                    default='similarity', help='panel选择方法')
    # parser.add_argument('--device', type=str, default='cuda', help='计算设备')
    # parser.add_argument('--model_path', type=str,
    #                    default='/home/zhangbo/workspace/aigc/SAM/vit_models/sam_vit_h_4b8939.pth',
    #                    help='SAM模型路径')
    
    args = parser.parse_args()
    
    print("📋 程序参数:")
    print(f"   📁 Panel目录: {args.panel_dir}")
    print(f"   📁 输出目录: {args.output_dir}")
    print(f"   📁 Mask目录: {args.mask_output_dir}")
    print(f"   📄 Forgery JSON: {args.forgery_json}")
    print(f"   🎯 选择方法: {args.method}")
    print(f"   🖥️  计算设备: {args.device}")
    print(f"   🤖 模型路径: {args.model_path}")
    print("=" * 80)
    
    # 检查输入目录
    print("🔍 检查输入目录...")
    if not os.path.exists(args.panel_dir):
        print(f"❌ 错误: panel目录不存在: {args.panel_dir}")
        return
    print("✅ 输入目录检查通过")
    
    # 创建处理器
    print("🔧 正在初始化处理器...")
    processor = PanelSplicingProcessor(device=args.device, model_path=args.model_path)
    print("✅ 处理器初始化完成")
    
    try:
        # 处理panel拼接
        result = processor.process_panel_splicing(
            panel_dir=args.panel_dir,
            output_dir=args.output_dir,
            mask_output_dir=args.mask_output_dir,
            method=args.method
        )
        
        # 更新forgery.json
        processor.update_forgery_json(args.forgery_json, result)
        
        print("=" * 80)
        print("🎉 所有处理完成!")
        print("=" * 80)
        print("📊 处理结果摘要:")
        print(f"   📄 原始Panel A: {result['original_panels']['panel_a']}")
        print(f"   📄 原始Panel B: {result['original_panels']['panel_b']}")
        print(f"   🎨 合成图像C: {result['synthetic_panels']['synthetic_c']}")
        print(f"   🎨 合成图像D: {result['synthetic_panels']['synthetic_d']}")
        print(f"   🎭 Mask A: {result['masks']['mask_a']}")
        print(f"   🎭 Mask B: {result['masks']['mask_b']}")
        print("=" * 80)
        
    except Exception as e:
        print("=" * 80)
        print(f"❌ 处理过程中出错: {e}")
        print("=" * 80)
        import traceback
        traceback.print_exc()
    
    finally:
        # 清理资源
        print("🧹 正在清理资源...")
        if processor.sam_model is not None:
            del processor.sam_model
        torch.cuda.empty_cache()
        gc.collect()
        print("✅ 资源清理完成")
        print("👋 程序结束")


if __name__ == "__main__":
    main()
